# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/2/6 10:39 PM
# @Function:
import re

import pandas as pd

from xinanjiang.util.Vector import Vector


class IO:
    def __init__(self):
        self.nrows = None
        self.ncols = None
        # 创建动态二维数据，用于存储各雨量站降雨量，mm
        self.m_P = Vector()
        # 创建动态二维数据，用于存储各蒸发站蒸发量，mm
        self.m_EM = Vector()
        self.m_Q = None

    def ReadFromFile(self, StrPath):
        # 定义要保存每行内容的变量列表
        _index = 0
        with open(StrPath + "/data/P.txt", 'r') as file:
            for line in file:
                line = line.strip()
                if _index == 0:
                    # 读入降雨数据的行数和列数，行数表示时段数，列数表示雨量站数
                    # print("第一行", line)
                    self.nrows = re.sub(r' .*', '', line)
                    self.ncols = re.sub(r'^.* ', '', line)
                    print("IO nrows", self.nrows)
                    print("IO ncols", self.ncols)
                else:
                    # 读取各雨量站降雨量数据，mm
                    lst = []
                    for item in line.split("\t"):
                        item = float(item.strip())
                        lst.append(item)
                    # print("数据", lst)
                    self.m_P.push_back(lst)
                _index = _index + 1
        # 输出结果
        print("m_P", self.m_P.data)
        print("m_P", self.m_P.size())

        _index = 0

        # 读入蒸发数据的行数和列数，行数表示时段数，列数表示蒸发站数
        with open(StrPath + "/data/EM.txt", 'r') as file:
            for line in file:
                line = line.strip()
                if _index == 0:
                    # 读入降雨数据的行数和列数，行数表示时段数，列数表示雨量站数
                    # print("第一行", line)
                    nrows = re.sub(r' .*', '', line)
                    ncols = re.sub(r'^.* ', '', line)
                    print("nrows", nrows)
                    print("ncols", ncols)
                else:
                    # 读取各雨量站降雨量数据，mm
                    lst = []
                    for item in line.split("\t"):
                        item = float(item.strip())
                        lst.append(item)
                    # print("数据", lst)
                    self.m_EM.push_back(lst)
                _index = _index + 1
        # 输出结果
        print("m_EM", self.m_EM.data)
        print("m_EM", self.m_EM.size())

    def WriteToFile(self, Path):
        pass


class Watershed:
    def __init__(self):
        self.m_areaSubWatershed = []
        self.m_numEvaporationStation = 0
        self.m_rateRainfallStation = Vector()
        self.m_numRainfallStation = 0
        self.m_numSubWatershed = 0
        self.m_rateEvaporationStation = Vector()
        self.m_P = Vector()
        self.m_EM = Vector()

    def ReadFromFile(self, StrPath):
        """
        读取流域分块信息
        Args:
            StrPath:

        Returns:

        """
        # 定义要保存每行内容的变量列表
        _index = 0

        # 读取雨量站点名称
        m_nameRainfallStation = []
        # 读取蒸发站点名称
        m_nameEvaporationStation = []

        # 创建动态二维数据，用于存储各雨量站降雨量，mm
        with open(StrPath + "/data/watershed.txt", 'r', encoding='utf-8') as file:
            for line in file:
                if _index == 0:
                    m_name = line
                    print("m_name", m_name)
                elif _index == 1:
                    m_area = line
                    print("m_area", m_area)
                elif _index == 2:
                    self.m_numRainfallStation = line
                    print("m_numRainfallStation", self.m_numRainfallStation)
                elif _index == 3:
                    self.m_numEvaporationStation = line
                    print("m_numEvaporationStation", self.m_numEvaporationStation)
                elif _index == 4:
                    # 读取单元流域面积
                    self.m_numSubWatershed = line
                    print("m_numSubWatershed", self.m_numSubWatershed)
                elif _index == 5:
                    # 读取单元流域面积
                    for item in line.split("\t"):
                        item = float(item.strip())
                        self.m_areaSubWatershed.append(item)
                    print("m_areaSubWatershed", self.m_areaSubWatershed)
                elif 5 < _index < int(self.m_numSubWatershed) + 6:
                    # 读取雨量站分配比例
                    # 创建子流域行 * 雨量站列的二维动态向量，并初始化为0
                    lst = []
                    for item in line.split("\t"):
                        item = float(item.strip())
                        lst.append(item)
                    self.m_rateRainfallStation.push_back(lst)
                elif 5 + int(self.m_numSubWatershed) < _index < int(self.m_numSubWatershed) * 2 + 6:
                    # 读取雨量站分配比例
                    # 创建子流域行 * 雨量站列的二维动态向量，并初始化为0
                    lst = []
                    for item in line.split("\t"):
                        item = float(item.strip())
                        lst.append(item)
                    self.m_rateEvaporationStation.push_back(lst)
                elif _index == int(self.m_numSubWatershed) * 2 + 6:
                    for item in line.split(" "):
                        item = item.strip()
                        m_nameRainfallStation.append(item)
                else:
                    for item in line.split(" "):
                        item = item.strip()
                        m_nameEvaporationStation.append(item)

                _index = _index + 1
        print("m_rateRainfallStation", self.m_rateRainfallStation.data)
        print("m_rateEvaporationStation", self.m_rateEvaporationStation.data)
        print("m_nameRainfallStation", m_nameRainfallStation)
        print("m_nameEvaporationStation", m_nameEvaporationStation)

    def calculate(self, io):
        """
        计算各单元流域各时段降雨和水面蒸发量
        Args:
            io:

        Returns:

        """
        # 记录条数，即时段数
        nrows = io.nrows
        # 单元流域个数
        ncols = self.m_numSubWatershed

        self.m_P.init_data(int(nrows), int(ncols))
        self.m_EM.init_data(int(nrows), int(ncols))
        for r in range(int(nrows)):
            for c in range(int(ncols)):
                for i in range(int(self.m_numRainfallStation)):
                    # 按比例计算单元流域降雨量
                    self.m_P.data[r][c] += \
                        io.m_P.data[r][i] * \
                        self.m_rateRainfallStation.data[c][i]
                for i in range(int(self.m_numEvaporationStation)):
                    # 计算各单元流域逐时段水面蒸发量，mm
                    self.m_EM.data[r][c] += \
                        io.m_EM.data[r][i] * \
                        self.m_rateEvaporationStation.data[c][i]
        # print("m_P", self.m_P.data)
        # print("m_P size()", self.m_P.size())
        print("m_EM", self.m_EM.data)
        print("m_EM size()", self.m_EM.size())

    def GetnW(self):
        return self.m_numSubWatershed

    def GetP(self, nt, nw):
        return self.m_P.data[nt][nw]

    def GetEM(self, nt, nw):
        return self.m_EM.data[nt][nw]

    def GetF(self, nw):
        return self.m_areaSubWatershed[nw]


if __name__ == '__main__':
    io = IO()
    io.ReadFromFile(".")
    Watershed = Watershed()
    Watershed.ReadFromFile(".")
    Watershed.calculate(io)
